package worker

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"

	sdebug "gitee.com/simplexyz/simplego/debug"
)

type Worker struct {
	started       atomic.Bool
	startWg       *sync.WaitGroup
	stopped       atomic.Bool
	stopWg        *sync.WaitGroup
	working       atomic.Bool
	ctx           context.Context
	cancelFunc    context.CancelFunc
	onStartedFunc func()
	onWorkingFunc func(ctx context.Context) bool
	onStopFunc    func()
	onErrorFunc   func(p any) error
}

func Create(startFunc func() error, wait bool, optionFuncs ...OptionFunc) (*Worker, error) {
	w := &Worker{}

	for _, f := range optionFuncs {
		f(w)
	}

	if w.ctx == nil {
		w.ctx, w.cancelFunc = context.WithCancel(context.Background())
	}

	if w.startWg == nil {
		w.startWg = &sync.WaitGroup{}
	}

	if w.stopWg == nil {
		w.stopWg = &sync.WaitGroup{}
	}

	if w.onErrorFunc == nil {
		w.onErrorFunc = func(p any) error {
			sdebug.WriteStackTraceFile(p, "")
			return p.(error)
		}
	}

	if err := w.start(startFunc, wait); err != nil {
		return nil, err
	}

	return w, nil
}

func (w *Worker) Working() bool {
	return w.working.Load()
}

func (w *Worker) work() (continued bool) {
	defer func() {
		if p := recover(); p != nil {
			if e := w.onErrorFunc(p); e != nil {
				continued = false
			} else {
				continued = true
			}
		}
	}()

	continued = w.onWorkingFunc(w.ctx)

	return
}

func (w *Worker) start(startFunc func() error, wait bool) error {
	if w.onWorkingFunc == nil {
		return ErrNoOnWorking
	}

	if !w.started.CompareAndSwap(false, true) {
		return ErrAlreadyStarted
	}

	if startFunc != nil {
		if err := startFunc(); err != nil {
			return fmt.Errorf("start worker fail, %w", err)
		}
	}

	if wait {
		w.startWg.Add(1)
	}

	w.stopWg.Add(1)

	go func() {
		if wait {
			w.startWg.Done()
		}

		if w.onStartedFunc != nil {
			w.onStartedFunc()
		}

		w.working.Store(true)

	workLoop:
		for {
			select {
			case <-w.ctx.Done():
				break workLoop
			default:
				if !w.work() {
					break workLoop
				}
			}
		}

		w.working.Store(false)

		if w.onStopFunc != nil {
			w.onStopFunc()
		}

		w.onStartedFunc = nil
		w.onWorkingFunc = nil
		w.onStopFunc = nil
		w.onErrorFunc = nil

		w.stopped.Store(true)
		w.stopWg.Done()
	}()

	if wait {
		w.startWg.Wait()
	}

	return nil
}

func (w *Worker) Stop(stopFunc func(), wait bool) {
	if !w.stopped.CompareAndSwap(false, true) {
		return
	}

	if stopFunc != nil {
		stopFunc()
	}

	if w.cancelFunc != nil {
		w.cancelFunc()
	}

	if wait {
		w.stopWg.Wait()
	}
}
